# -*- coding:utf-8 -*-

import os, sys, re
import numpy as np
import pandas as pd
import pandas.json as _json

from filecache.models import *
from datakeeper.utils import *
# from metlib.kits import *
from metlib.misc.misc import isseq
from metlib.io.misc import *
from windsim import to_wws

try:
    data_cache_bucket = FileCacheBucket.objects.get(name='data_cache')
except FileCacheBucket.DoesNotExist:
    data_cache_bucket = FileCacheBucket(name='data_cache')
    data_cache_bucket.save()

UNUSED_PARAMS = {
    'uri_info', 'cache',
}

def get_uri_with_params(uri, params):
    param_str = '&'.join(['%s=%s' % (key, params[key]) for key in sorted(params.keys()) if key not in UNUSED_PARAMS])
    return uri + '?' + param_str

def save_data_cache(data, uri_with_params, bucket=data_cache_bucket):
    uri = uri_with_params + '&PKL'
    try:
        cache = bucket.get(uri, when_not_exist=('create',))
        savepickle(cache.filepath, data)
        new_size = DU(cache.filepath)
        cache.update(size=new_size)
    except Exception as e:
        try:
            cache.delete()
        except Exception as e:
            pass

def load_data_cache(uri_with_params):
    uri = uri_with_params + '&PKL'
    try:
        cache = FileCache.objects.get(pk=uri)
        fname = cache.filepath
        data = loadpickle(fname)
        return data
    except FileCache.DoesNotExist:
        return None


def gen_datapack(packuri, parts, *args, **kwargs):
    """使用一批DataUnit生成DataPack.
    :param parts: (UnitUri, Unit)组成的列表.
    :return: DataPack
    """
    res = {
        "type": "datapack",
        "uri": packuri,
        "tags": {},
        "contents": {}
    }
    for parturi, part in parts:
        if part['type'] == 'dataunit':
            res['tags'][parturi] = parturi
            res['contents'][parturi] = part
        elif part['type'] == 'datapack':
            for subtag, suburi in part['tags'].iteritems():
                res['tags'][suburi] = suburi   # use suburi as tagname
                res['contents'][suburi] = part['contetns'][subtag]
    return res


def get_zipped_data(dataset, uris, res_uri, request=None, *args, **kwargs):
    # TODO add common support, both for this function and for parsing subcontents.
    res = {
        "type": "datapack",
        "uri": res_uri,
        "tags": {},
        "contents": {}
    }
    for uri in uris:
        try:
            data = dataset.get_data(uri, request=request, *args, **kwargs)
            if data['type'] == 'dataunit':
                res['tags'][uri] = uri
                res['contents'][uri] = data
            elif data['type'] == 'datapack':
                for subtag, suburi in data['tags'].iteritems():
                    res['tags'][suburi] = suburi
                    res['contents'][suburi] = data['contents'][subtag]
        except Exception as e:
            continue

    return res


def convert2dataunit(info={}):
    res = {
        "type": "dataunit",
        "info": {},
        "data": {},
    }
    for key, value in info:
        if key not in res:
            res[key] = value
        else:
            res[key].update(value)
    return res


def pd_ts2dataunit(ts, info={}, equal_interval=True):
    pass

def pd_df2dataunit(df, info={}, equal_interval=True):
    pass

def dts_values2dataunit(dts, values, info={}, equal_interval=True):
    pass

def amend_dataunit(unit, info={}):
    unit['type'] = 'dataunit'
    for key, value in info.iteritems():
        if key not in unit:
            unit[key] = value
        else:
            unit[key].update(value)
    return unit


def amend_datapack(pack, info={}):
    pack['type'] = 'datapack'
    for tag, uri in pack['tags'].iteritems():
        unit = pack['contents'][tag]
        amend_dataunit(unit, info)
    return pack


def get_no_access_dataunit(uri):
    return {
        "type": "dataunit",
        "uri": uri,
        "error": "Access Forbidden",
    }


def get_error_dataunit(uri):
    return {
        "type": "dataunit",
        "uri": uri,
        "error": "Get Data Error"
    }


# 装饰器之cache_data, 自动提供data_cache
def cache_data(valid_params=None):
    def outer_wrapper(get_data_func):
        def modified_func(instance, uri, request=None, **kwargs):
            cache = kwargs.get('cache', 'rw')
            if valid_params is None or callable(valid_params):
                cache_kwargs = kwargs
            elif isseq(valid_params):
                cache_kwargs = {k:kwargs[k] for k in valid_params if k in kwargs}
            else:
                datatype = kwargs.get('datatype', None)
                if datatype is not None and datatype in valid_params:
                    this_valid_params = valid_params.get(datatype)
                    cache_kwargs = {k:kwargs[k] for k in this_valid_params if k in kwargs}
                else:
                    cache_kwargs = kwargs

            uri_with_params = get_uri_with_params(uri, cache_kwargs)
            if cache in ['rw', 'r']:
                data = load_data_cache(uri_with_params)
                if data is not None:
                    return data

            data = get_data_func(instance, uri, request=request, **kwargs)

            if cache in ['rw', 'w']:
                save_data_cache(data, uri_with_params)
            return data

        return modified_func

    if callable(valid_params):  # no args given to cache_data, so use valid_params as the given func.
        return outer_wrapper(valid_params)
    else:
        return outer_wrapper


# 装饰器之pack_data, 自动提供datapack机制
def pack_data(get_data_func):
    def modified_func(instance, uri, request=None, **kwargs):
        datatype = kwargs.get('datatype', 'datapack')
        if datatype != 'datapack':
            return get_data_func(instance, uri, request=request, **kwargs)

        modified_kwargs = kwargs.copy()
        modified_kwargs.update({
            'datatype': 'dataunit',
            'uri_info': None
        })
        parts = []
        suburi_list = instance.unzip_uri(uri)
        for suburi in suburi_list:
            subdata = get_data_func(instance, suburi, request=request, **modified_kwargs)
            parts.append((suburi, subdata))
        res = gen_datapack(uri, parts)
        return res

    return modified_func


# 装饰器之权限检测
# 用在PT上时需谨慎
def perm_data(get_data_func):
    def modified_func(instance, uri, request=None, **kwargs):
        kwargs.pop('sample', None)
        try:
            # has_perm = check_data_perm(instance, uri, request=request, **kwargs)
            if request is None:
                if instance.is_public():
                    min_uri = instance.minimize_uri(uri)
                    if instance.is_free(min_uri):
                        has_perm = True
                    else:
                        has_perm = False
                else:
                    has_perm = 'NoAccess'
            else:
                has_perm = user_owns_data(request.user, uri)

            if has_perm == 'NoAccess':
                raise DataNoAccess

            data = get_data_func(instance, uri, request=request, sample=not has_perm, **kwargs)
            return data
        except DataNoAccess:
            return get_no_access_dataunit(uri)
        except GetDataError:
            return get_error_dataunit(uri)

    return modified_func


# 装饰器之provide_figinfo, 用在RY的get_data上
def provide_figinfo(func=None, **kws):
    def outer_wrapper(get_data_func):
        def modified_func(instance, uri, request=None, *args, **kwargs):
            datatype = kwargs.get('datatype', 'figinfo')
            if datatype == 'figinfo':
                uri_info = kwargs.get('uri_info', None)
                if uri_info is None:
                    uri_info = instance.parse_uri(uri)
                varname = uri_info['varname']
                url_para = 'format=<%= format %>' + \
                           '&projection=<%= projection %>' + \
                           '&pitem=<%= pitem %>' + \
                           '&vmin=<%= vmin %>&vmax=<%= vmax %>&cmap=<%= cmap %>'

                edited_kwargs = kwargs.copy()
                edited_kwargs.update({
                    'datatype': 'dataunit',
                })
                dataunit = get_data_func(instance, uri, request=request, *args, **edited_kwargs)
                data = dataunit['data']['values']
                suggest_range = dataunit['data'].get('suggest_range', (None, None))
                valid_data = data[np.isfinite(data)]
                fit_vmin = np.percentile(valid_data, 5)
                fit_vmax = np.percentile(valid_data, 95)
                fit_ext = (fit_vmax - fit_vmin) / 4.0
                fit_vmin -= fit_ext
                fit_vmax += fit_ext
                fit_range = (fit_vmin, fit_vmax)
                fix_range = list(suggest_range)
                if fix_range[0] is None:
                    fix_range[0] = fit_vmin
                if fix_range[1] is None:
                    fix_range[1] = fit_vmax

                res = {
                    "type": "datapack",
                    "uri": uri,
                    "tags": {
                        "data": uri
                    },
                    "common": {
                        "info": {}
                    },
                    "contents": {
                        "data": {
                            "type": "dataunit",
                            "uri": uri,
                            "info": {
                                "varname": varname,
                                "type": "rect_figure",
                                "sample": dataunit['info']['sample'],
                            },
                            "coords": dataunit.get('coords', {}),
                            "data": {
                                "url_para": url_para,
                                "units": dataunit['data'].get('units', ''),
                                "suggest_range": suggest_range,
                                "suggest_cmap": kws.get('suggest_cmap_dict', {}).get(varname, None),
                                "fit_range": fit_range,
                                "fix_range": fix_range,
                            }
                        }
                    }
                }
                return res
            else:
                return get_data_func(instance, uri, request=request, **kwargs)
        return modified_func
    if callable(func):
        return outer_wrapper(func)
    else:
        return outer_wrapper


# 装饰器之csv支持, 目前用在PT上, 要求get_data()能支持datatype='dataframe'.
def support_csv(get_file_func):
    def modified_func(instance, uri, dest, request=None, *args, **kwargs):
        format = kwargs.get('format', 'pack')
        if format == 'csv':
            kwargs['datatype'] = 'dataframe'
            kwargs.pop('format', None)
            data = instance.get_data(uri, request=request, *args, **kwargs)
            if len(data) == 0:
                return {'status': 'NoData'}
            data.to_csv(dest, na_rep='NaN', float_format="%.4f")
            info = {
                'filetype': 'csv',
                'format': 'csv',
                'status': 'OK'
            }
            return info
        else:
            res = get_file_func(instance, uri, dest, request=request, *args, **kwargs)
            if res:
                return res
            else:
                return {'status': 'NoData'}
    return modified_func


# 装饰器之wws(WindSim格式), 目前用在PY的风玫瑰上
def support_wws(get_file_func):
    def modified_func(instance, uri, dest, request=None, *args, **kwargs):
        format = kwargs.get('format', 'pack')
        if format == 'wws':
            kwargs['datatype'] = 'rawrose'
            kwargs.pop('format', None)
            data = instance.get_data(uri, request=request, *args, **kwargs)
            if not data:
                return {'status': 'NoData'}
            to_wws(data, dest)
            info = {
                'filetype': 'wws',
                'format': 'wws',
                'status': 'OK'
            }
            return info
        else:
            res = get_file_func(instance, uri, dest, request=request, *args, **kwargs)
            if res:
                return res
            else:
                return {'status': 'NoData'}
    return modified_func


def support_packunitjson(get_file_func):
    def modified_func(instance, uri, dest, request=None, *args, **kwargs):
        format = kwargs.get('format', 'pack')
        if format == 'json':
            kwargs.pop('format', None)
            if ',' in uri or '*' in uri:
                kwargs['datatype'] = 'datapack'
            else:
                kwargs['datatype'] = 'dataunit'
        elif format == 'pack':
            kwargs.pop('format', None)
            kwargs['datatype'] = 'datapack'
        else:
            res = get_file_func(instance, uri, dest, request=request, *args, **kwargs)
            if res:
                return res
            else:
                return {'status': 'NoData'}

        data = instance.get_data(uri, request=request, *args, **kwargs)
        if len(data) == 0:
            return {'status': 'NoData'}

        info = {}
        with open(dest, 'w') as outf:
            outf.write(_json.dumps(data, date_unit='s', double_precision=4))
        info['filetype'] = kwargs['datatype']
        info['format'] = 'json'
        info['status'] = 'OK'
        return info
    return modified_func


# 装饰器之防hack (这个过滤层可能放在datakeeper那边更合适), 用在顶层的get_file()上.
def protect_hack(get_file_func):
    def modified_func(instance, uri, dest, request=None, *args, **kwargs):
        kwargs.pop('sample', None)

        res = get_file_func(instance, uri, dest, request=request, *args, **kwargs)
        return res

    return modified_func
